#include "taskManager.hpp"
#include <chrono>
#include <iostream>
#include <thread>
#include <vector>
#include <cmath>
#include <sstream>
#include "log.h"

using namespace std;

QueueSet* TaskManager::_qs = nullptr;
vector<TaskManager*> TaskManager::workers;
int TaskManager::worker_num = 0;
TaskManager::Terminator* TaskManager::terminator = new TaskManager::Terminator();

static int parallel_threads = 64;
static int termination_wait_millisecond = 10;

void TaskManager::start(TaskManager* tm) {
    tm->work();
}

void TaskManager::initialize(Graph &graph){
    worker_num = parallel_threads;
    _qs = new QueueSet(worker_num);

    int gcTaskNumber = ceil((float)graph.getGCRootNumber() / worker_num);
    for(int i = 0; i < worker_num; i++) {
        TaskManager* tm = new TaskManager(i, _qs->get_queue(i));
        tm->pushGCRoot(graph, i * gcTaskNumber, gcTaskNumber);
        workers.push_back(tm); 
        thread* t = new thread(start, tm);
        tm->set_thread(t);
    }

}

void TaskManager::pushGCRoot(Graph& graph, int start, int taskNumber) {
    for(int i = 0; i < taskNumber && start + i < graph.getGCRootNumber(); i++) {
        push(new Task(graph.getGCRootNode(start + i)));
    }
}

void TaskManager::end() {
    for(int i = 0; i < worker_num; i++) {
        workers[i]->get_thread()->join();
        std::stringstream ss;
        ss << "Thread " << workers[i]->_id << ": process nums are " << workers[i]->getProcessNum();
        LOG(workers[i]->output, ss.str())
    }
    for(int i = 0; i < worker_num; i++) {
        delete workers[i];
    }

    delete _qs;
}

bool TaskManager::steal(int id, Task *&t) {
    t = _qs->steal(id, this);
    return t != nullptr;
}

TaskManager::~TaskManager() {
    output.close();
}

void TaskManager::work() {
    while(true) {
        drain_stack();
        Task* t;
        std::stringstream ss;
        ss << "Thread " << _id << ": " << "steal";
        LOG(output, ss.str());
        if(TaskManager::steal(_id, t)) {
            push(t);
        } else {
            std::stringstream ss;
            ss << "Thread " << _id << ": " << "steal fail ";
            LOG(output, ss.str());
            if(terminator->offer_termination()) {
                break;
            }
        }
    }
}

void TaskManager::push(Task* task) {
    task_stack->push(task);
}

void TaskManager::drain_stack() {
    Task* t = nullptr;
    while((t = task_stack->pop()) != nullptr) {
        t->do_task(this);
    }
}

int TaskManager::getId() {
    return _id;
}

std::ofstream& TaskManager::getOutput() {
    return output;
}

void TaskManager::addProcessNum() {
    processNum++;
}

int TaskManager::getProcessNum() {
    return processNum;
}

bool TaskManager::Terminator::offer_termination() {
    if(worker_num == 1) {
        _offered_termination = 1;
        return true;
    }

    std::unique_lock<std::mutex> lk(_lock);

    //TODO: need check
    _offered_termination++;

    if(_offered_termination >= worker_num) {
        prepare_for_return();
        return true;
    }

    bool should_terminate = cv.wait_for(lk,
            std::chrono::milliseconds(termination_wait_millisecond),
            [this]{ return _offered_termination == worker_num; });

    if(_offered_termination >= worker_num) {
        should_terminate = true; 
    }

    if (should_terminate) {
        prepare_for_return();
        return true;
    } else {
        _offered_termination--;
        return false;
    }

    return true;
}

void TaskManager::Terminator::prepare_for_return() {
    std::thread::id this_thread = std::this_thread::get_id();
    //if (this_thread == _spin_master->get_id()) {
    //    _spin_master = nullptr; 
    //}
    cv.notify_all();
}
